"""
数据集中各种种类的数量占比
1、为了尽可能使数据集对各个种类的合理分配学习资源，请保证各种类数量接近
2、输出文件列表中的总红色数量、蓝色数量、黄色数量及其占比
"""
############需要修改###################
# 定义文件夹路径列表，需要统计的文件夹（可以放测试集、训练集、验证集等）
folder_paths = [
'/home/hw/dataset/cone_total_new/labels',

]
#####################################

import os

# 定义统计函数
def count_lines(file_path):
    counts = {'0': 0, '1': 0, '2': 0}
    with open(file_path, 'r') as file:
        for line in file:
            if line.startswith('0'):
                counts['0'] += 1
            elif line.startswith('1'):
                counts['1'] += 1
            elif line.startswith('2'):
                counts['2'] += 1
    return counts

# 统计每个文件夹中的txt文件
total_num = 0
total_counts = {'0': 0, '1': 0, '2': 0}
for folder_path in folder_paths:
    # 获取文件夹中所有txt文件
    txt_files = [f for f in os.listdir(folder_path) if f.endswith('.txt')]
    
    # 统计每个txt文件中的行数
    for txt_file in txt_files:
        file_path = os.path.join(folder_path, txt_file)
        counts = count_lines(file_path)
        for key in counts:
            total_counts[key] += counts[key]

# 输出统计结果
for key, value in total_counts.items():
    total_num += value
print("类别统计结果:")
print(f"总数 {total_num}")

for key, value in total_counts.items():
    print(f"类别 {key}: {value}\t占比：{round(float(value)/float(total_num)*100,2)}%")

